import os
import json
import logging
import argparse  # 添加命令行参数解析支持
import sys  # 添加系统模块
from dotenv import load_dotenv
from openai import OpenAI
import importlib

# 配置日志系统
logging.basicConfig(level=logging.INFO, 
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    handlers=[
                        logging.FileHandler("deepseek_api.log", encoding="utf-8"),
                        logging.StreamHandler()
                    ])
logger = logging.getLogger(__name__)

# 将系统信息输出到日志
logger.info(f"系统编码信息：默认编码={sys.getdefaultencoding()}, 控制台编码={sys.stdout.encoding}")

# 检查OpenAI库版本
openai_version = importlib.import_module("openai").__version__
logger.info(f"当前使用的OpenAI库版本: {openai_version}")

# 直接从.env文件读取API密钥，不依赖系统环境变量
def load_env_file(env_path='.env'):
    """
    直接从.env文件读取变量，不依赖系统环境变量
    """
    env_vars = {}
    try:
        if os.path.exists(env_path):
            with open(env_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith('#'):
                        continue
                    key, value = line.split('=', 1)
                    env_vars[key.strip()] = value.strip()
            logger.info(f"成功从{env_path}加载环境变量")
        else:
            logger.warning(f"环境变量文件{env_path}不存在")
    except Exception as e:
        logger.error(f"读取环境变量文件出错: {str(e)}")
    
    return env_vars

# 读取环境变量文件
env_vars = load_env_file()

# 获取API密钥
DEEPSEEK_API_KEY = env_vars.get('DEEPSEEK_API_KEY')
DEEPSEEK_API_URL = env_vars.get('DEEPSEEK_API_URL', 'https://api.deepseek.com')

logger.info(f"DeepSeek API密钥: {'已配置' if DEEPSEEK_API_KEY else '未配置'}")
logger.info(f"DeepSeek API URL: {DEEPSEEK_API_URL}")

class DeepseekAPI:
    def __init__(self):
        if not DEEPSEEK_API_KEY:
            logger.warning('未找到Deepseek API密钥，请在.env文件中设置DEEPSEEK_API_KEY')
        
        # 根据OpenAI库版本决定初始化参数
        try:
            # 检查OpenAI版本以决定初始化参数
            openai_version_parts = openai_version.split('.')
            major_version = int(openai_version_parts[0])
            minor_version = int(openai_version_parts[1])
            
            # OpenAI 1.6.0不支持proxies参数
            if openai_version == "1.6.0" or (major_version == 1 and minor_version < 7):
                logger.info(f"使用OpenAI v{openai_version}初始化客户端（不含proxies参数）")
                self.client = OpenAI(
                    api_key=DEEPSEEK_API_KEY,
                    base_url=DEEPSEEK_API_URL  # 官方文档推荐的基础URL
                )
            else:
                # 版本1.7.0及以上可能需要使用代理
                logger.info(f"使用OpenAI v{openai_version}初始化客户端（支持proxies参数）")
                http_proxy = os.environ.get('http_proxy')
                https_proxy = os.environ.get('https_proxy')
                
                if http_proxy or https_proxy:
                    proxies = {
                        "http": http_proxy,
                        "https": https_proxy
                    }
                    self.client = OpenAI(
                        api_key=DEEPSEEK_API_KEY,
                        base_url=DEEPSEEK_API_URL,
                        proxies=proxies
                    )
                else:
                    self.client = OpenAI(
                        api_key=DEEPSEEK_API_KEY,
                        base_url=DEEPSEEK_API_URL
                    )
        except Exception as e:
            logger.error(f"初始化OpenAI客户端时出错: {str(e)}")
            logger.warning("使用基础参数初始化OpenAI客户端")
            self.client = OpenAI(
                api_key=DEEPSEEK_API_KEY,
                base_url=DEEPSEEK_API_URL
            )
    
    def chat_completion(self, messages, model="deepseek-chat", max_tokens=2048, stream=False):
        """执行聊天完成请求
        
        Args:
            messages: 对话历史消息列表
            model: 使用的模型，可选值: 'deepseek-chat'(V3模型) 或 'deepseek-reasoner'(R1推理模型)
            max_tokens: 最大输出token数，范围1-8192
            stream: 是否使用流式输出
            
        Returns:
            如果stream=False，返回完整响应
            如果stream=True，返回流式响应生成器
        """
        try:
            # 检查API密钥
            if not DEEPSEEK_API_KEY or DEEPSEEK_API_KEY == 'your_deepseek_key_here':
                logger.warning('Deepseek API密钥未正确配置，返回模拟响应')
                if not stream:
                    return self._mock_completion(messages)
                else:
                    return self._mock_stream_completion(messages)
            
            # 验证模型参数
            if model not in ["deepseek-chat", "deepseek-reasoner"]:
                logger.warning(f'不支持的模型: {model}，将使用默认模型 deepseek-chat')
                model = "deepseek-chat"
            
            # 验证max_tokens参数
            if max_tokens < 1 or max_tokens > 8192:
                logger.warning(f'max_tokens超出范围: {max_tokens}，将使用默认值2048')
                max_tokens = 2048
                
            logger.info(f'调用DeepSeek API，模型: {model}, 流式输出: {stream}')
            
            # 根据模型类型设置参数
            params = {
                "model": model,
                "messages": messages,
                "max_tokens": max_tokens,
                "stream": stream
            }
            
            # 对于deepseek-chat模型，可以设置temperature和top_p
            if model == "deepseek-chat":
                params["temperature"] = 0.7
                params["top_p"] = 0.9
            
            # 使用openai客户端调用API
            response = self.client.chat.completions.create(**params)
            
            logger.info('API调用成功')
            return response
            
        except Exception as e:
            logger.error(f'调用DeepSeek API失败: {str(e)}')
            raise

    def review_document(self, content, prompt_template=None):
        """使用deepseek-reasoner模型审核文档内容
        
        Args:
            content: 需要审核的文档内容
            prompt_template: 提示词模板，如果为None则使用默认模板
            
        Returns:
            审核结果和思考过程
        """
        try:
            if not prompt_template:
                prompt_template = """##角色:你是一位经验丰富的技术文档审核专家。
##任务:请仔细审核以下文档内容，评估其质量、完整性和专业性，并提供改进建议。

{content}

请从以下几个方面进行评审：
1. 文档结构与组织：评估文档的整体结构是否清晰、逻辑性强
2. 技术内容准确性：检查技术描述是否准确，术语使用是否正确
3. 完整性：评估文档是否包含所有必要的信息
4. 可读性：考察文档的语言表达、格式和排版
5. 改进建议：提供具体的改进建议

请首先进行深入思考，然后给出评审结果。评审结果应简明扼要，重点突出需要改进的地方。
"""
            
            prompt = prompt_template.format(content=content)
            
            messages = [
                {"role": "user", "content": prompt}
            ]
            
            logger.info('正在使用R1推理模型审核文档...')
            
            # 使用推理模型进行文档审核
            response = self.chat_completion(
                messages=messages,
                model="deepseek-reasoner",  # 使用R1推理模型
                max_tokens=4096,
                stream=False
            )
            
            # 提取审核结果和思考过程（reasoning_content）
            if hasattr(response, 'choices') and len(response.choices) > 0:
                content = response.choices[0].message.content
                # 检查是否有reasoning_content字段（推理过程）
                reasoning_content = getattr(response.choices[0].message, 'reasoning_content', None)
                
                logger.info(f'审核完成，结果长度: {len(content)} 字符')
                if reasoning_content:
                    logger.info(f'包含推理过程，长度: {len(reasoning_content)} 字符')
                
                return {
                    "result": content,
                    "thinking": reasoning_content or "分析过程已包含在结果中"
                }
            else:
                logger.warning('审核失败，无法获取结果')
                return {
                    "result": "审核失败，无法获取结果",
                    "thinking": "API响应格式异常"
                }
                
        except Exception as e:
            logger.error(f'文档审核失败: {str(e)}')
            return {
                "result": f"审核过程发生错误: {str(e)}",
                "thinking": "API调用失败"
            }
    
    def _mock_completion(self, messages):
        """模拟API响应（用于API密钥未配置时）"""
        user_message = messages[-1]['content'] if messages else "空消息"
        logger.warning(f'使用模拟响应代替实际API调用，用户消息: {user_message[:30]}...')
        
        # 创建一个模拟的响应对象
        class MockResponse:
            def __init__(self, content, reasoning_content=None):
                self.choices = [MockChoice(content, reasoning_content)]
                
        class MockChoice:
            def __init__(self, content, reasoning_content=None):
                self.message = MockMessage(content, reasoning_content)
                
        class MockMessage:
            def __init__(self, content, reasoning_content=None):
                self.content = content
                self.reasoning_content = reasoning_content
        
        # 如果消息包含特定关键词，生成带有reasoning_content的模拟响应
        if "审核" in user_message or "评审" in user_message or "分析" in user_message:
            reasoning = "我需要仔细分析这个文档...\n\n首先，我会检查文档结构...\n\n然后，我会评估技术内容的准确性...\n\n最后，我会提出改进建议..."
            content = "文档分析结果：\n1. 文档结构：整体结构清晰\n2. 内容准确性：基本准确，但有几处需要修正\n3. 改进建议：建议补充更多示例"
            return MockResponse(content, reasoning)
        else:
            return MockResponse(f"这是一个模拟响应。\n您的消息是：{user_message}\n模拟回答：这是一个模拟回答。")

    def _mock_stream_completion(self, messages):
        """生成模拟的流式完成响应，返回一个生成器"""
        user_message = messages[-1]['content'] if messages and messages[-1]['role'] == 'user' else ''
        
        # 模拟思维链内容 (reasoning_content)
        mock_reasoning = f"思考分析问题：'{user_message}'\n\n我需要理解用户的问题并给出合适的回答。让我分析一下。\n\n此问题涉及到...\n\n经过分析我认为..."
        
        # 模拟最终答案
        mock_content = f"您好！关于'{user_message}'的问题，我的回答是...请问还有其他问题吗？"
        
        # 将响应分解为小块
        reasoning_chunks = [mock_reasoning[i:i+5] for i in range(0, len(mock_reasoning), 5)]
        content_chunks = [mock_content[i:i+5] for i in range(0, len(mock_content), 5)]
        
        # 创建模拟的流式响应类
        class MockDelta:
            def __init__(self, content=None, reasoning_content=None):
                self.content = content
                self.reasoning_content = reasoning_content
                
        class MockChoice:
            def __init__(self, delta):
                self.delta = delta
                
        class MockChunk:
            def __init__(self, choices):
                self.choices = choices
        
        # 首先流式输出思维链内容
        for chunk in reasoning_chunks:
            delta = MockDelta(reasoning_content=chunk)
            choice = MockChoice(delta)
            yield MockChunk([choice])
        
        # 然后流式输出最终内容
        for chunk in content_chunks:
            delta = MockDelta(content=chunk)
            choice = MockChoice(delta)
            yield MockChunk([choice])

# 添加命令行参数解析
def parse_arguments():
    parser = argparse.ArgumentParser(description='DeepSeek API调用工具')
    parser.add_argument('--mode', choices=['chat', 'review'], required=True, 
                        help='调用模式：chat-聊天模式，review-文档审核模式')
    parser.add_argument('--input', required=True, help='输入JSON文件路径')
    parser.add_argument('--output', required=True, help='输出JSON文件路径')
    parser.add_argument('--stream', action='store_true', help='是否使用流式输出（仅聊天模式）')
    return parser.parse_args()

# 主函数
def main():
    # 解析命令行参数
    args = parse_arguments()
    
    try:
        logger.info(f'开始处理，模式: {args.mode}, 输入文件: {args.input}, 输出文件: {args.output}')
        
        # 初始化API客户端
        api = DeepseekAPI()
        
        # 读取输入文件
        with open(args.input, 'r', encoding='utf-8') as f:
            input_data = json.load(f)
        
        result = None
        
        # 根据模式执行不同操作
        if args.mode == 'chat':
            # 聊天模式
            if not input_data.get('messages'):
                raise ValueError('输入文件缺少messages字段')
            
            if args.stream:
                # 流式输出处理
                logger.info('开始执行流式聊天完成请求...')
                response_stream = api.chat_completion(
                    messages=input_data['messages'],
                    model=input_data.get('model', 'deepseek-chat'),
                    max_tokens=input_data.get('max_tokens', 2048),
                    stream=True
                )
                
                # 收集流式响应
                content = ""
                reasoning_content = ""
                
                for chunk in response_stream:
                    if hasattr(chunk.choices[0].delta, 'reasoning_content') and chunk.choices[0].delta.reasoning_content:
                        reasoning_content += chunk.choices[0].delta.reasoning_content
                    elif hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
                        content += chunk.choices[0].delta.content
                
                result = {
                    'content': content,
                    'reasoning_content': reasoning_content if reasoning_content else None
                }
            else:
                logger.info('开始执行聊天完成请求...')
                response = api.chat_completion(
                    messages=input_data['messages'],
                    model=input_data.get('model', 'deepseek-chat'),
                    max_tokens=input_data.get('max_tokens', 2048)
                )
                
                if hasattr(response, 'choices') and len(response.choices) > 0:
                    result = {
                        'content': response.choices[0].message.content,
                        'reasoning_content': getattr(response.choices[0].message, 'reasoning_content', None)
                    }
                else:
                    result = {'error': '无法获取响应结果'}
        
        elif args.mode == 'review':
            # 文档审核模式
            if 'content' not in input_data:
                raise ValueError('输入文件缺少content字段')
            
            logger.info('开始执行文档审核...')
            logger.info(f'文档长度: {len(input_data["content"])} 字符')
            
            review_result = api.review_document(
                content=input_data['content'],
                prompt_template=input_data.get('prompt_template')
            )
            
            result = review_result
        
        # 写入输出文件
        logger.info(f'写入结果到输出文件: {args.output}')
        with open(args.output, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        
        logger.info(f'成功完成 {args.mode} 任务')
        
    except Exception as e:
        error_msg = f'处理失败: {str(e)}'
        logger.error(error_msg)
        # 写入错误信息到输出文件
        try:
            with open(args.output, 'w', encoding='utf-8') as f:
                json.dump({'error': str(e)}, f, ensure_ascii=False, indent=2)
        except Exception as write_error:
            logger.error(f'写入错误信息失败: {str(write_error)}')
        
        sys.exit(1)

# 为直接运行提供入口
if __name__ == "__main__":
    main()
